/*
    This code is meant to analyze the Portable Executable (PE) file format.
        [+] PE Header [ e_lfanew (0x3C) = Contains Offset to NT Header (0x80) ]
        [+] COFF File Header peHeader, machine, NumberofSections, TimeDateStamp, SizeOfOptionalHeader
        [+] Section Header
        [+] Export Directory
        [+] RVA & RVAOffset

    Author: @5mukx
*/

use pe_files::{analyze_rich_header, print_dos_header};
use winapi::um::synchapi::WaitForSingleObject;
use winapi::um::{
    memoryapi::VirtualProtect, minwinbase::LPTHREAD_START_ROUTINE,
    processthreadsapi::CreateRemoteThread,
};

mod reflect;
use std::ptr::null_mut;
use winapi::um::winnt::PAGE_EXECUTE_READ;
mod pe_files;
// use pe_parser::{PeFile, PeResult};

// importing from another image

// use crate::reflect::*;

const MAX_PATH: usize = 260;

fn find_export_directory_info(
    pointer_to_raw_data: usize,
    virtual_address_offset: usize,
    pe_analyzer: &[u8],
) -> usize {
    println!("INSIDE FIND EXPORT DIRECTORY !");
    println!(
        "         [-] [0x{:04x}] [exportFlags] : 0x{:02X}{:02X}{:02X}{:02X}",
        pointer_to_raw_data,
        pe_analyzer[pointer_to_raw_data + 3],
        pe_analyzer[pointer_to_raw_data + 2],
        pe_analyzer[pointer_to_raw_data + 1],
        pe_analyzer[pointer_to_raw_data]
    );

    let time_date_stamp_offset = pointer_to_raw_data + 0x4;
    println!(
        "         [-] [0x{:04x}] [Time/DateStamp] : 0x{:02X}{:02X}{:02X}{:02X}",
        time_date_stamp_offset,
        pe_analyzer[time_date_stamp_offset + 3],
        pe_analyzer[time_date_stamp_offset + 2],
        pe_analyzer[time_date_stamp_offset + 1],
        pe_analyzer[time_date_stamp_offset]
    );

    let major_version_offset = time_date_stamp_offset + 0x4;
    println!(
        "         [-] [0x{:04x}] [majorVersion] : 0x{:02X}{:02X}",
        major_version_offset,
        pe_analyzer[major_version_offset + 1],
        pe_analyzer[major_version_offset]
    );

    let minor_version_offset = major_version_offset + 0x2;
    println!(
        "         [-] [0x{:04x}] [minorVersion] : 0x{:02X}{:02X}",
        minor_version_offset,
        pe_analyzer[minor_version_offset + 1],
        pe_analyzer[minor_version_offset]
    );

    let name_rva_offset = minor_version_offset + 0x2;
    println!(
        "         [-] [0x{:04x}] [nameRVA] : 0x{:02X}{:02X}{:02X}{:02X} (RVA to PE name)",
        name_rva_offset,
        pe_analyzer[name_rva_offset + 3],
        pe_analyzer[name_rva_offset + 2],
        pe_analyzer[name_rva_offset + 1],
        pe_analyzer[name_rva_offset]
    );

    let ordinal_base_offset = name_rva_offset + 0x4;
    println!(
        "         [-] [0x{:04x}] [ordinalBase] : 0x{:02X}{:02X}{:02X}{:02X}",
        ordinal_base_offset,
        pe_analyzer[ordinal_base_offset + 3],
        pe_analyzer[ordinal_base_offset + 2],
        pe_analyzer[ordinal_base_offset + 1],
        pe_analyzer[ordinal_base_offset]
    );

    let address_table_entries_offset = ordinal_base_offset + 0x4;
    println!(
        "         [-] [0x{:04x}] [addressTableEntries] : 0x{:02X}{:02X}{:02X}{:02X} (Count of functions in Export Address Table)",
        address_table_entries_offset,
        pe_analyzer[address_table_entries_offset + 3],
        pe_analyzer[address_table_entries_offset + 2],
        pe_analyzer[address_table_entries_offset + 1],
        pe_analyzer[address_table_entries_offset]
    );

    let number_of_name_pointers_offset = address_table_entries_offset + 0x4;
    println!(
        "         [-] [0x{:04x}] [numberOfNamePointers] : 0x{:02X}{:02X}{:02X}{:02X} (Count of entries in the name pointer table/ordinal table)",
        number_of_name_pointers_offset,
        pe_analyzer[number_of_name_pointers_offset + 3],
        pe_analyzer[number_of_name_pointers_offset + 2],
        pe_analyzer[number_of_name_pointers_offset + 1],
        pe_analyzer[number_of_name_pointers_offset]
    );

    let export_address_table_rva_offset = number_of_name_pointers_offset + 0x4;
    println!(
        "         [-] [0x{:04x}] [exportAddressTableRVA] : 0x{:02X}{:02X}{:02X}{:02X} (RVA of the Export Address Table)",
        export_address_table_rva_offset,
        pe_analyzer[export_address_table_rva_offset + 3],
        pe_analyzer[export_address_table_rva_offset + 2],
        pe_analyzer[export_address_table_rva_offset + 1],
        pe_analyzer[export_address_table_rva_offset]
    );

    let name_pointer_rva_offset = export_address_table_rva_offset + 0x4;
    println!(
        "         [-] [0x{:04x}] [namePointerRVA] : 0x{:02X}{:02X}{:02X}{:02X} (RVA of the Export Name Pointer Table)",
        name_pointer_rva_offset,
        pe_analyzer[name_pointer_rva_offset + 3],
        pe_analyzer[name_pointer_rva_offset + 2],
        pe_analyzer[name_pointer_rva_offset + 1],
        pe_analyzer[name_pointer_rva_offset]
    );

    let ordinal_table_rva_offset = name_pointer_rva_offset + 0x4;
    println!(
        "         [-] [0x{:04x}] [ordinalTableRVA] : 0x{:02X}{:02X}{:02X}{:02X} (RVA of the Ordinal Table)",
        ordinal_table_rva_offset,
        pe_analyzer[ordinal_table_rva_offset + 3],
        pe_analyzer[ordinal_table_rva_offset + 2],
        pe_analyzer[ordinal_table_rva_offset + 1],
        pe_analyzer[ordinal_table_rva_offset]
    );

    let export_name_pointer_rva = ((pe_analyzer[name_pointer_rva_offset + 3] as usize) << 24)
        | ((pe_analyzer[name_pointer_rva_offset + 2] as usize) << 16)
        | ((pe_analyzer[name_pointer_rva_offset + 1] as usize) << 8)
        | (pe_analyzer[name_pointer_rva_offset] as usize);

    let edata_virtual_address = ((pe_analyzer[virtual_address_offset + 3] as usize) << 24)
        | ((pe_analyzer[virtual_address_offset + 2] as usize) << 16)
        | ((pe_analyzer[virtual_address_offset + 1] as usize) << 8)
        | (pe_analyzer[virtual_address_offset] as usize);

    let export_name_pointer_file_offset =
        (export_name_pointer_rva - edata_virtual_address) + pointer_to_raw_data;

    println!(
        "         [-] [0x{:04x}] [exportNamePointerRVA] : 0x{:02X}{:02X}{:02X}{:02X}",
        export_name_pointer_file_offset,
        pe_analyzer[export_name_pointer_file_offset + 3],
        pe_analyzer[export_name_pointer_file_offset + 2],
        pe_analyzer[export_name_pointer_file_offset + 1],
        pe_analyzer[export_name_pointer_file_offset]
    );

    let symbol_name_rva = ((pe_analyzer[export_name_pointer_file_offset + 3] as usize) << 24)
        | ((pe_analyzer[export_name_pointer_file_offset + 2] as usize) << 16)
        | ((pe_analyzer[export_name_pointer_file_offset + 1] as usize) << 8)
        | (pe_analyzer[export_name_pointer_file_offset] as usize);

    let symbol_file_offset = (symbol_name_rva - edata_virtual_address) + pointer_to_raw_data;
    let mut symbol_name = [0u8; MAX_PATH];

    for i in 0..MAX_PATH {
        if pe_analyzer[symbol_file_offset + i] == 0 {
            break;
        }
        symbol_name[i] = pe_analyzer[symbol_file_offset + i];
    }

    let symbol_name_str = String::from_utf8_lossy(&symbol_name[..]);
    println!(
        "         [-] [0x{:04x}] [symbolName] : {}",
        symbol_file_offset, symbol_name_str
    );

    let export_address_table_rva = ((pe_analyzer[export_address_table_rva_offset + 3] as usize)
        << 24)
        | ((pe_analyzer[export_address_table_rva_offset + 2] as usize) << 16)
        | ((pe_analyzer[export_address_table_rva_offset + 1] as usize) << 8)
        | (pe_analyzer[export_address_table_rva_offset] as usize);

    println!(
        "         Calculation = (0x{:X} - 0x{:X}) + 0x{:X}",
        export_address_table_rva, edata_virtual_address, pointer_to_raw_data
    );

    let symbol_rva_offset =
        (export_address_table_rva - edata_virtual_address) + pointer_to_raw_data;
    let symbol_rva = ((pe_analyzer[symbol_rva_offset + 3] as usize) << 24)
        | ((pe_analyzer[symbol_rva_offset + 2] as usize) << 16)
        | ((pe_analyzer[symbol_rva_offset + 1] as usize) << 8)
        | (pe_analyzer[symbol_rva_offset] as usize);
    println!(
        "         [-] [0x{:04x}] [symbolRVA] : 0x{:08X}",
        symbol_rva_offset, symbol_rva
    );
    symbol_rva
}

fn find_section_headers(
    first_section_header_offset: usize,
    no_of_section: usize,
    pe_analyzer: &[u8],
) {
    println!(
        "\n [Sections headers start at: 0x{:04x}]",
        first_section_header_offset
    );

    // table header
    println!(
        "{:<8} {:<10} {:<12} {:<15} {:<14} {:<17} {:<20} {:<20} {:<19} {:<15}",
        "Offset",
        "Name",
        "VirtualSize",
        "VirtualAddress",
        "SizeOfRawData",
        "PointerToRawData",
        "PointerToRelocations",
        "PointerToLinenumbers",
        "NumberOfLinenumbers",
        "Characteristics"
    );
    println!("{}", "-".repeat(145));

    let mut next_section_header_offset = first_section_header_offset;

    for _ in 0..no_of_section {
        // Extract the section name (8 bytes)
        let header_name = String::from_utf8_lossy(
            &pe_analyzer[next_section_header_offset..next_section_header_offset + 8],
        )
        .trim_end_matches('\0')
        .to_string();

        // Extract VirtualSize (4 bytes)
        let virtual_size_offset = next_section_header_offset + 0x8;
        let virtual_size = u32::from_le_bytes([
            pe_analyzer[virtual_size_offset],
            pe_analyzer[virtual_size_offset + 1],
            pe_analyzer[virtual_size_offset + 2],
            pe_analyzer[virtual_size_offset + 3],
        ]);

        // Extract VirtualAddress (4 bytes)
        let virtual_address_offset = virtual_size_offset + 0x4;
        let virtual_address = u32::from_le_bytes([
            pe_analyzer[virtual_address_offset],
            pe_analyzer[virtual_address_offset + 1],
            pe_analyzer[virtual_address_offset + 2],
            pe_analyzer[virtual_address_offset + 3],
        ]);

        // Extract SizeOfRawData (4 bytes)
        let size_of_raw_data_offset = virtual_address_offset + 0x4;
        let size_of_raw_data = u32::from_le_bytes([
            pe_analyzer[size_of_raw_data_offset],
            pe_analyzer[size_of_raw_data_offset + 1],
            pe_analyzer[size_of_raw_data_offset + 2],
            pe_analyzer[size_of_raw_data_offset + 3],
        ]);

        // Extract PointerToRawData (4 bytes)
        let pointer_to_raw_data_offset = size_of_raw_data_offset + 0x4;
        let pointer_to_raw_data = u32::from_le_bytes([
            pe_analyzer[pointer_to_raw_data_offset],
            pe_analyzer[pointer_to_raw_data_offset + 1],
            pe_analyzer[pointer_to_raw_data_offset + 2],
            pe_analyzer[pointer_to_raw_data_offset + 3],
        ]);

        // Extract PointerToRelocations (4 bytes)
        let pointer_to_relocations_offset = pointer_to_raw_data_offset + 0x4;
        let pointer_to_relocations = u32::from_le_bytes([
            pe_analyzer[pointer_to_relocations_offset],
            pe_analyzer[pointer_to_relocations_offset + 1],
            pe_analyzer[pointer_to_relocations_offset + 2],
            pe_analyzer[pointer_to_relocations_offset + 3],
        ]);

        // Extract PointerToLinenumbers (4 bytes)
        let pointer_to_linenumbers_offset = pointer_to_relocations_offset + 0x4;
        let pointer_to_linenumbers = u32::from_le_bytes([
            pe_analyzer[pointer_to_linenumbers_offset],
            pe_analyzer[pointer_to_linenumbers_offset + 1],
            pe_analyzer[pointer_to_linenumbers_offset + 2],
            pe_analyzer[pointer_to_linenumbers_offset + 3],
        ]);

        // Extract NumberOfLinenumbers (4 bytes)
        let number_of_linenumbers_offset = pointer_to_linenumbers_offset + 0x4;
        let number_of_linenumbers = u32::from_le_bytes([
            pe_analyzer[number_of_linenumbers_offset],
            pe_analyzer[number_of_linenumbers_offset + 1],
            pe_analyzer[number_of_linenumbers_offset + 2],
            pe_analyzer[number_of_linenumbers_offset + 3],
        ]);

        // Extract Characteristics (4 bytes)
        let characteristics_offset = number_of_linenumbers_offset + 0x4;
        let characteristics = u32::from_le_bytes([
            pe_analyzer[characteristics_offset],
            pe_analyzer[characteristics_offset + 1],
            pe_analyzer[characteristics_offset + 2],
            pe_analyzer[characteristics_offset + 3],
        ]);

        // Print the row for this section
        println!(
            "{:<8x} {:<10} {:<12x} {:<15x} {:<14x} {:<17x} {:<20x} {:<20x} {:<19x} {:<15x}",
            next_section_header_offset,
            header_name,
            virtual_size,
            virtual_address,
            size_of_raw_data,
            pointer_to_raw_data,
            pointer_to_relocations,
            pointer_to_linenumbers,
            number_of_linenumbers,
            characteristics
        );

        // Existing logic for .edata section
        if header_name.contains(".edata") {
            let pointer_to_raw_data = u32::from_le_bytes([
                pe_analyzer[pointer_to_raw_data_offset],
                pe_analyzer[pointer_to_raw_data_offset + 1],
                pe_analyzer[pointer_to_raw_data_offset + 2],
                pe_analyzer[pointer_to_raw_data_offset + 3],
            ]);
            let symbol_rva = find_export_directory_info(
                pointer_to_raw_data as usize,
                virtual_address_offset,
                pe_analyzer,
            );

            let mut temp_section_header_offset = first_section_header_offset;
            println!(
                "temp_section_header_offset value: {:#?}",
                temp_section_header_offset
            );

            for _ in 0..11 {
                let section_virtual_address_offset = first_section_header_offset + 0xC;

                let section_virtual_address = u32::from_le_bytes([
                    pe_analyzer[section_virtual_address_offset],
                    pe_analyzer[section_virtual_address_offset + 1],
                    pe_analyzer[section_virtual_address_offset + 2],
                    pe_analyzer[section_virtual_address_offset + 3],
                ]);
                let section_size_of_raw_data_offset = section_virtual_address_offset + 0x4;

                let section_size_of_raw_data = u32::from_le_bytes([
                    pe_analyzer[section_size_of_raw_data_offset],
                    pe_analyzer[section_size_of_raw_data_offset + 1],
                    pe_analyzer[section_size_of_raw_data_offset + 2],
                    pe_analyzer[section_size_of_raw_data_offset + 3],
                ]);
                let section_pointer_to_raw_data_offset = section_size_of_raw_data_offset + 0x4;

                let section_pointer_to_raw_data = u32::from_le_bytes([
                    pe_analyzer[section_pointer_to_raw_data_offset],
                    pe_analyzer[section_pointer_to_raw_data_offset + 1],
                    pe_analyzer[section_pointer_to_raw_data_offset + 2],
                    pe_analyzer[section_pointer_to_raw_data_offset + 3],
                ]);

                if symbol_rva > section_virtual_address as usize
                    && symbol_rva < (section_virtual_address + section_size_of_raw_data) as usize
                {
                    let symbol_file_offset = (symbol_rva - section_virtual_address as usize)
                        + section_pointer_to_raw_data as usize;
                    println!(
                        "     [*] [0x{:04x}] [symbolFileOffset] : 0x{:08X}",
                        symbol_rva, symbol_file_offset
                    );

                    let mut pe_analyzer_executable_buffer: Vec<u8> =
                        vec![0u8; pe_analyzer.len()];
                    pe_analyzer_executable_buffer
                        .copy_from_slice(&pe_analyzer[..pe_analyzer.len()]);

                    println!(
                        "------------------------ boxreflectDllExectuableBuffer: 0x{:x}",
                        pe_analyzer_executable_buffer.as_ptr() as usize
                    );

                    for i in 0..0x84 {
                        if i == 0x3c {
                            continue;
                        }
                        pe_analyzer_executable_buffer[i] = 0;
                    }

                    pe_analyzer_executable_buffer[0x80] = 0x23;
                    pe_analyzer_executable_buffer[0x81] = 0x12;

                    let fl_old_protect: u32 = 0;

                    unsafe {
                        VirtualProtect(
                            pe_analyzer_executable_buffer.as_mut_ptr() as *mut _,
                            pe_analyzer.len(),
                            PAGE_EXECUTE_READ,
                            fl_old_protect.clone() as *mut u32,
                        );
                    }

                    let symbol_executable_address: LPTHREAD_START_ROUTINE = unsafe {
                        std::mem::transmute(
                            (pe_analyzer_executable_buffer.as_ptr() as usize + symbol_file_offset)
                                as usize,
                        )
                    };

                    let mut lp_thread_id = 0;
                    let h_thread = unsafe {
                        CreateRemoteThread(
                            null_mut(),
                            null_mut(),
                            1024 * 1024,
                            symbol_executable_address,
                            null_mut(),
                            0,
                            &mut lp_thread_id,
                        )
                    };
                    unsafe {
                        WaitForSingleObject(h_thread, 0xFFFFFFFF);
                    }
                    break;
                }
                temp_section_header_offset += 0x28;
            }
        }
        next_section_header_offset += 0x28;
    }
}


fn main() {
    let initial_offset: usize = 0x3c;

/*
    If you want to analyze an shellcode dll etc.., you can follow the import method , call from the function
    If you want to analyze an file , you can use include_bytes to read from file.
*/

    //                      ----- IMPORT METHOD -----
    let args: Vec<String> = std::env::args().collect();

    if args.len() < 2 {
        println!("[i] Please Enter PE File to parse ...");
        return;
    }

    let pe_analyzer = std::fs::read(&args[1]).expect("Failed to read the PE file");


    let pe_header_offset: usize = (pe_analyzer[initial_offset + 2] as usize) << 16
        | (pe_analyzer[initial_offset + 1] as usize) << 8
        | pe_analyzer[initial_offset] as usize;

    println!(
        " [0x{:x}] [peHeader offset] : 0x{:x}",
        initial_offset, pe_header_offset
    );

    println!(
        " [0x{:x}] [peHeader] : {}{}",
        pe_header_offset,
        pe_analyzer[pe_header_offset] as char,
        pe_analyzer[pe_header_offset + 1] as char
    );

    // calculate machineTypeOffset
    let machine_type_offset: usize = pe_header_offset + 4;
    println!(
        " [0x{:x}] [machineType] : x{:x}{:x}",
        machine_type_offset,
        pe_analyzer[machine_type_offset + 1],
        pe_analyzer[machine_type_offset]
    );

    // calculate noOfSectionsOffset and noOfSections
    let no_of_sections_offset: usize = machine_type_offset + 0x2;
    let no_of_sections: u16 = ((pe_analyzer[no_of_sections_offset + 1] as u16) << 8)
        | (pe_analyzer[no_of_sections_offset] as u16);

    println!(
        " [0x{:x}] [noOfSections] : ({}) 0x{:x}",
        no_of_sections_offset, no_of_sections, no_of_sections
    );

    // calculate timeDateStampOffset
    let time_date_stamp_offset: usize = no_of_sections_offset + 0x2;
    println!(
        " [0x{:x}] [timeDateStamp] : 0x{:x}{:x}{:x}{:x}",
        time_date_stamp_offset,
        pe_analyzer[time_date_stamp_offset + 3],
        pe_analyzer[time_date_stamp_offset + 2],
        pe_analyzer[time_date_stamp_offset + 1],
        pe_analyzer[time_date_stamp_offset]
    );

    // calculate sizeOfOptionalHeaderOffset and sizeOfOptionalHeader
    let size_of_optional_header_offset: usize = time_date_stamp_offset + 0x4 + 0x4 + 0x4;
    let size_of_optional_header: u16 = ((pe_analyzer[size_of_optional_header_offset + 1] as u16)
        << 8)
        | (pe_analyzer[size_of_optional_header_offset] as u16);
    println!(
        " [0x{:x}] [sizeOfOptionalHeader] : 0x{:x}",
        size_of_optional_header_offset, size_of_optional_header
    );

    // calculate firstSectionHeaderOffset
    let first_section_header_offset =
        size_of_optional_header_offset + 0x2 + 0x2 + size_of_optional_header as usize;

    println!();
    
    // Include the HEADERS  -> 
    // Analyze DOS Header Here !

    print_dos_header(&pe_analyzer);
    
    println!();

    // Calculate RichHeader
    // rich header adv implementaion ! 
    
    analyze_rich_header(&pe_analyzer);

    println!("\n\n");

    find_section_headers(
        first_section_header_offset,
        no_of_sections as usize,
        &pe_analyzer,
    );
}
